##########################################################################################

library(data.table)
library(dplyr)
library(ggplot2)
library(ggpubr)
library(optparse)
library(tidyr)

##########################################################################################

option_list <- list(
  make_option(c("--info_file"), type = "character"),
  make_option(c("--out_path"), type = "character")
)

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

if(1!=1){
  info_file <- "/public/home/xxf2019/20220915_gastric_multiple/dna_combinePublic/public_ref/combine/MutationInfo.combine.addMolecularSubType.rmMIX.tsv"
}

info_file <- opt$info_file
out_path <- opt$out_path

dir.create(out_path , recursive = T)

###########################################################################################

dat_info <- data.frame(fread(info_file))

###########################################################################################
## 构建All
dat_info_all <- dat_info
dat_info_all$Class <- "IGC + DGC"

dat_info <- rbind(dat_info,dat_info_all)

###########################################################################################
## 计算和画图函数
plofFun <- function( dat_info_use = dat_info_use , base = base ){
  dat_info_use$use_col <- dat_info_use[[base]]

  ## 按照患者计算
  dat_plot <- dat_info_use %>%
    group_by(Class,Molecular.subtype,use_col) %>%
    summarise(count=n()) %>% 
    group_by(Class , use_col) %>%
    mutate(count_all=sum(count),
           ratio=count/count_all)
  dat_plot$Molecular.subtype[dat_plot$Molecular.subtype=="MSI"] <- "MSI/POLE"

  dat_plot$p <- ""

  for(i in unique(dat_plot$Class)){
    dat_p <- subset(dat_plot,Class==i)
    dat_p <- dat_p[,c("Molecular.subtype","use_col","count")]
    matrix_data <- dat_p %>%
      group_by(use_col) %>%
      spread(Molecular.subtype, count,fill = 0)
    if(nrow(matrix_data)==2){
      dat_plot$p[dat_plot$Class==i][1] <- fisher.test( matrix_data[,-1] )$p.value
    } else {
      dat_plot$p[dat_plot$Class==i] <- ""
    }
  }

  trans <- function(num){
    up <- floor(log10(num))
    down <- round(num*10^(-up),2)
    text <- paste("P == ",down," %*% 10","^",up)
    return(text)
  }
  dat_plot$p <- as.numeric(dat_plot$p)
  dat_plot$value_text <- round(dat_plot$ratio , 2) * 100 
  dat_plot$p_text <- ifelse( dat_plot$p < 0.01 , trans(dat_plot$p) , paste0( "P == " , round(as.numeric(dat_plot$p) , 2) )  )
  dat_plot$p_text <- ifelse( is.na(dat_plot$p_text) , "" , dat_plot$p_text  )

  images_name <- paste0(out_path , "/molecular_subtype_baseinfo." , base , ".tsv")
  write.table( dat_plot , images_name , row.names = F , sep = "\t" , quote = F )

  ##############################################
  ## 画图
  col <- c(
    rgb(204,116,92,alpha=255,maxColorValue=255),
    rgb(62,62,93,alpha=255,maxColorValue=255),
    rgb(229,198,143,alpha=255,maxColorValue=255)
  )

  names(col) <- c("CIN","GS","MSI/POLE")
  dat_plot$Molecular.subtype <- factor(dat_plot$Molecular.subtype,levels=c("CIN","GS","MSI/POLE"))
  dat_plot$Class <- factor(dat_plot$Class,levels=c("IGC + DGC","IGC","DGC"))
  dat_plot$use_col <- factor(dat_plot$use_col,levels=unique(dat_plot$use_col ))
  dat_plot$use_col_n <- paste0( dat_plot$use_col , "\n(" , dat_plot$count_all , ")")

  if(base == "Alcohol"){
    dat_plot$use_col_n <- factor( dat_plot$use_col_n , levels = unique(dat_plot$use_col_n )[order(unique(dat_plot$use_col_n ) , decreasing=F )] )
  }else{
    dat_plot$use_col_n <- factor( dat_plot$use_col_n , levels = unique(dat_plot$use_col_n )[order(unique(dat_plot$use_col_n ) , decreasing=T )] )
  }

  p <- ggplot(dat_plot , mapping = aes(use_col_n , ratio , fill = Molecular.subtype)) +
    geom_bar(position = "stack", stat = "identity") + 
    theme_bw() +
    facet_grid(.~Class,space='free_x',scales='free_x') +
    theme(strip.text.x = element_text(size = 7 , face = 'bold'))+
    theme(panel.grid=element_blank())+
    labs(y = 'Proportions (%)') +
    geom_text(aes(label=p_text , y = 1.05 ,x=1.5),size=3,parse = TRUE)+
    geom_text(aes(label=value_text) , position=position_stack(vjust =0.5) , size=4 , color="white")+
    xlab("") +
    scale_fill_manual(values=col) +
    theme(
      title =element_text(size=5, face='bold'),
      strip.text.x = element_text(size = 10, colour = "black",face='bold') ,
      axis.text.x = element_text(size = 8,color="black",face='bold') ,
      axis.text.y = element_text(size = 8,color="black",face='bold') ,
      axis.title.y = element_text(size = 10,color="black",face='bold') ,
      #axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1 , size = 8,color="black",face='bold'),
      legend.title = element_blank(),
      legend.text = element_text(size = 8),
      legend.position = "bottom",
      axis.title =element_text(size = 8),axis.text =element_text(size = 7, color = 'black')
    )


  images_name <- paste0(out_path , "/molecular_subtype_baseinfo." , base , ".pdf")
  ggsave( images_name , p , width = 6 , height = 5 )

}

###########################################################################################
## 分三种基线计算
base_type <- c("Tobacco" , "Alcohol" , "HP")
for( base in base_type ){
  dat_info_use <- dat_info[dat_info[[base]]!="unknown",]
  plofFun( dat_info_use = dat_info_use , base = base )
}
